package com.example.spider.provider.multidownload.handler;

import com.example.spider.provider.multidownload.task.DownloadFilePartTask;
import com.example.spider.provider.multidownload.task.ParseDownloadFileTask;
import com.example.spider.task.SpiderResultHandler;
import com.example.spider.task.SpiderTaskQueue;
import com.example.spider.task.limit.LimitScope;
import com.example.spider.task.limit.MaxCurrentRequestLimit;
import lombok.extern.slf4j.Slf4j;
import org.shoulder.batch.progress.BatchProgressCache;
import org.shoulder.core.exception.CommonErrorCodeEnum;
import org.shoulder.core.util.AssertUtils;
import org.shoulder.core.util.JsonUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.HttpHeaders;
import org.springframework.stereotype.Component;
import org.springframework.util.unit.DataSize;

import java.time.Duration;
import java.util.Map;
import java.util.Optional;
import java.util.stream.Collectors;

/**
 * 处理列表查询结果
 *
 * @author lym
 */
@Slf4j
@Component
public class ParseDownloadFileTaskHandler implements SpiderResultHandler<ParseDownloadFileTask, byte[]> {

    @Autowired
    private BatchProgressCache progressCache;

    private long partSize = DataSize.ofBytes(1024).toBytes();

    /**
     * 分析列表结果，创建下载详情任务
     */
    @Override
    public void handle(ParseDownloadFileTask task, byte[] result) {
        Map<String, String> respHeaders = task.getResponseHeaders();
        Map<String, String> formatKeyRespHeaders = respHeaders.entrySet().stream()
                .collect(Collectors.toMap(e -> e.getKey().toLowerCase(), Map.Entry::getValue));
        //binary/octet-stream
//        String contentType = formatKeyRespHeaders.get(HttpHeaders.CONTENT_TYPE.toLowerCase());
        long contentLength = Optional.ofNullable(formatKeyRespHeaders.get(HttpHeaders.CONTENT_LENGTH.toLowerCase()))
                .map(Long::parseLong)
                .orElse(-1L);
        AssertUtils.isTrue(contentLength > 1L, CommonErrorCodeEnum.RPC_UNKNOWN, "unSupport such contentLength. taskId=" + task.getId());

        task.setFileSize(contentLength);

        String eTag = formatKeyRespHeaders.get(HttpHeaders.ETAG.toLowerCase());
        task.setEtag(eTag);

        String acceptRanges = formatKeyRespHeaders.get(HttpHeaders.ACCEPT_RANGES.toLowerCase());
        if (!"bytes".equals(acceptRanges)) {
            // 支持分段下载 + 断点续传
            task.setTotalPart(1);
        } else {
            // todo config 1kb
            task.setTotalPart((int) (contentLength / partSize) + 1);
        }

        // todo config limit
        MaxCurrentRequestLimit limitStrategy = new MaxCurrentRequestLimit("download:" + task.getId() + ":" + task.getEtag(),
                Duration.ofSeconds(1), LimitScope.CUSTOMER, 12);
        for (int i = 0; i < task.getTotalPart(); i++) {
            DownloadFilePartTask partTask = new DownloadFilePartTask();
            partTask.setUrl(task.getUrl());
            partTask.setIndex(i);
            partTask.setStart(i * partSize);
            partTask.setEnd(i == task.getTotalPart() - 1 ? task.getFileSize() : (i + 1) * partSize - 1);
            partTask.setParent(task);
            partTask.setLimitStrategy(limitStrategy);
            partTask.setHeaders(JsonUtils.parseObject(JsonUtils.toJson(task.getHeaders()), Map.class, String.class, String.class));
            partTask.getHeaders().put("Range", "bytes=%d-%d".formatted(partTask.getStart(), partTask.getEnd()));
            SpiderTaskQueue.putTask(partTask);
        }
    }

}
